import os
from time import sleep

import openai
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.llms import OpenAI
from langchain_community.vectorstores.chroma import Chroma
from langchain_core.example_selectors import SemanticSimilarityExampleSelector
from langchain_core.prompts import PromptTemplate, FewShotPromptTemplate

os.environ["OPENAI_API_KEY"]=os.environ["OPENAI_API_KEY_ZHIHU"]
os.environ["OPENAI_API_BASE"]=os.environ["OPENAI_API_BASE_ZHIHU"]
llm = OpenAI(model_name= "gpt-3.5-turbo",temperature=0.9)

#示例三
examples = [
    {"input":"happy","output":"sad"},
    {"input":"tall","output":"short"},
    {"input":"energetic","output":"lethargic"},
    {"input":"sunny","output":"gloomy"},
    {"input":"windy","output":"calm"},
]

example_template = """Input: {input}\nOutput: {output}"""

example_prompt = PromptTemplate(input_variables=["input","output"],template=example_template)

prefix = """Give the antonym of every input"""
suffix = """Input: {adjective}\nOutput:"""

example_selector = SemanticSimilarityExampleSelector.from_examples(examples=examples,embeddings=OpenAIEmbeddings(),vectorstore_cls=Chroma,k=1)

few_shot_prompt_template = FewShotPromptTemplate(prefix=prefix,suffix=suffix,example_selector=example_selector,input_variables=["adjective"],example_prompt=example_prompt)
query = "What defines a network engineer"
print(few_shot_prompt_template.format(adjective="worried"))
sleep(30)
print(few_shot_prompt_template.format(adjective="cloudy"))
sleep(30)
print(few_shot_prompt_template.format(adjective="tired"))
